// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/testing/gtest_utils.h>
#include <pollux/testing/functions/prestosql/function_base_test.h>

namespace kumo::pollux::functions {
namespace {

class FindFirstTest : public functions::test::FunctionBaseTest {
 protected:
  void verify(
      const std::string& expression,
      const RowVectorPtr& input,
      const VectorPtr& expected) {
    SCOPED_TRACE(expression);
    auto result = evaluate(expression, input);
    pollux::test::assertEqualVectors(expected, result);
  }
};

TEST_F(FindFirstTest, basic) {
  auto data = make_row_vector({
      make_array_vector_from_json<int32_t>({
          "[1, 2, 3, 4]",
          "null",
          "[]",
          "[null, 10, 20, null, 30, null, 40]",
      }),
  });

  // find_first: x > 0.
  VectorPtr expected =
      make_nullable_flat_vector<int32_t>({1, std::nullopt, std::nullopt, 10});
  verify("find_first(c0, x -> (x > 0))", data, expected);

  expected =
      make_nullable_flat_vector<int64_t>({1, std::nullopt, std::nullopt, 2});
  verify("find_first_index(c0, x -> (x > 0))", data, expected);

  // find_first: x > 10.
  expected = make_nullable_flat_vector<int32_t>(
      {std::nullopt, std::nullopt, std::nullopt, 20});
  verify("find_first(c0, x -> (x > 10))", data, expected);

  expected = make_nullable_flat_vector<int64_t>(
      {std::nullopt, std::nullopt, std::nullopt, 3});
  verify("find_first_index(c0, x -> (x > 10))", data, expected);

  // find_first: x < 0.
  expected = make_null_constant(TypeKind::INTEGER, 4);
  verify("find_first(c0, x -> (x < 0))", data, expected);

  expected = make_null_constant(TypeKind::BIGINT, 4);
  verify("find_first_index(c0, x -> (x < 0))", data, expected);

  // find_first x > 2 starting with 2nd element.
  expected =
      make_nullable_flat_vector<int32_t>({3, std::nullopt, std::nullopt, 10});
  verify("find_first(c0, 2, x -> (x > 2))", data, expected);

  expected =
      make_nullable_flat_vector<int64_t>({3, std::nullopt, std::nullopt, 2});
  verify("find_first_index(c0, 2, x -> (x > 2))", data, expected);

  // find_first x > 2 starting with 5-th element.
  expected = make_nullable_flat_vector<int32_t>(
      {std::nullopt, std::nullopt, std::nullopt, 30});
  verify("find_first(c0, 5, x -> (x > 2))", data, expected);

  expected = make_nullable_flat_vector<int64_t>(
      {std::nullopt, std::nullopt, std::nullopt, 5});
  verify("find_first_index(c0, 5, x -> (x > 2))", data, expected);

  // first_first x > 0 from the end of the array.
  expected =
      make_nullable_flat_vector<int32_t>({4, std::nullopt, std::nullopt, 40});
  verify("find_first(c0, -1, x -> (x > 0))", data, expected);

  expected =
      make_nullable_flat_vector<int64_t>({4, std::nullopt, std::nullopt, 7});
  verify("find_first_index(c0, -1, x -> (x > 0))", data, expected);

  // first_first x > 0 from the 2-nd to last element of the array.
  expected =
      make_nullable_flat_vector<int32_t>({3, std::nullopt, std::nullopt, 30});
  verify("find_first(c0, -2, x -> (x > 0))", data, expected);

  expected =
      make_nullable_flat_vector<int64_t>({3, std::nullopt, std::nullopt, 5});
  verify("find_first_index(c0, -2, x -> (x > 0))", data, expected);

  expected = make_null_constant(TypeKind::INTEGER, 4);
  verify("find_first(c0, cast(null as INTEGER), x -> (x > 0))", data, expected);

  expected = make_null_constant(TypeKind::BIGINT, 4);
  verify(
      "find_first_index(c0, cast(null as INTEGER), x -> (x > 0))",
      data,
      expected);
}

TEST_F(FindFirstTest, firstMatchIsNull) {
  auto data = make_row_vector({
      make_array_vector_from_json<int32_t>({
          "[1, null, 2]",
      }),
  });

  POLLUX_ASSERT_THROW(
      evaluate("find_first(c0, x -> (x is null))", data),
      "find_first found NULL as the first match");
}

TEST_F(FindFirstTest, predicateFailures) {
  auto data = make_row_vector({make_array_vector_from_json<int32_t>({
      "[1, 2, 3, 0]",
      "[-1, 3, 0, 5]",
      "[5, 6, 7, 0]",
  })});

  POLLUX_ASSERT_THROW(
      evaluate("find_first(c0, x -> (10 / x > 2))", data), "division by zero");
  POLLUX_ASSERT_THROW(
      evaluate("find_first(c0, 2, x -> (10 / x > 2))", data),
      "division by zero");
  POLLUX_ASSERT_THROW(
      evaluate("find_first_index(c0, x -> (10 / x > 2))", data),
      "division by zero");
  POLLUX_ASSERT_THROW(
      evaluate("find_first_index(c0, 2, x -> (10 / x > 2))", data),
      "division by zero");

  VectorPtr expected = make_nullable_flat_vector<int32_t>({
      1,
      3,
      std::nullopt,
  });

  verify("find_first(c0, x -> (try(10 / x) > 2))", data, expected);
  verify("try(find_first(c0, x -> (10 / x > 2)))", data, expected);

  expected = make_nullable_flat_vector<int32_t>({
      2,
      3,
      std::nullopt,
  });
  verify("find_first(c0, -3, x -> (10 / x > 2))", data, expected);

  expected = make_nullable_flat_vector<int64_t>({
      1,
      2,
      std::nullopt,
  });
  verify("find_first_index(c0, x -> (try(10 / x) > 2))", data, expected);
  verify("try(find_first_index(c0, x -> (10 / x > 2)))", data, expected);

  expected = make_nullable_flat_vector<int64_t>({
      2,
      2,
      std::nullopt,
  });
  verify("find_first_index(c0, -3, x -> (10 / x > 2))", data, expected);
}

TEST_F(FindFirstTest, invalidIndex) {
  auto data = make_row_vector({
      make_array_vector_from_json<int32_t>({
          "[1, 2, 3]",
          "[-1, 2]",
          "[-2, -3, -4]",
          "[-5, -6]",
      }),
      make_flat_vector<int32_t>({2, 0, 0, 1}),
  });

  // Index 0 is not valid. Expect an error.
  POLLUX_ASSERT_THROW(
      evaluate("find_first(c0, c1, x -> (x > 0))", data),
      "SQL array indices start at 1. Got 0.");

  // There are 2 rows with invalid index. Mark array argument in one of these
  // rows as NULL. Still expect the other row to trigger an error.
  data->childAt(0)->set_null(1, true);
  POLLUX_ASSERT_THROW(
      evaluate("find_first(c0, c1, x -> (x > 0))", data),
      "SQL array indices start at 1. Got 0.");

  // Mark array argument in the other row as NULL. Expect no errors.
  data->childAt(0)->set_null(2, true);
  auto expected = make_nullable_flat_vector<int32_t>(
      {2, std::nullopt, std::nullopt, std::nullopt});
  verify("find_first(c0, c1, x -> (x > 0))", data, expected);

  // All null or empty arrays.
  data = make_row_vector({
      make_array_vector_from_json<int32_t>({
          "[]",
          "null",
          "[]",
          "null",
      }),
      make_flat_vector<int32_t>({2, 0, 0, 1}),
  });

  // Index 0 is not valid. Expect an error in the 3rd row.
  POLLUX_ASSERT_THROW(
      evaluate("find_first(c0, c1, x -> (x > 0))", data),
      "SQL array indices start at 1. Got 0.");

  // Mark 3rd row null. Expect no error.
  data->childAt(1)->set_null(2, true);
  expected = make_all_null_flat_vector<int32_t>(4);
  verify("find_first(c0, c1, x -> (x > 0))", data, expected);
}

// Verify that null arrays with non-zero offsets/sizes are processed correctly.
TEST_F(FindFirstTest, nulls) {
  auto data = make_row_vector({
      make_array_vector_from_json<int32_t>({
          "[1, 2, 3]",
          "[-1, 2]",
          "[-2, -3, -4]",
          "[-5, -6]",
      }),
  });

  // find_first: x > 0.
  VectorPtr expected =
      make_nullable_flat_vector<int32_t>({1, 2, std::nullopt, std::nullopt});
  verify("find_first(c0, x -> (x > 0))", data, expected);

  // Mark [-1, 2] array as null. Expect null result.
  data->childAt(0)->set_null(1, true);
  expected = make_nullable_flat_vector<int32_t>(
      {1, std::nullopt, std::nullopt, std::nullopt});
  verify("find_first(c0, x -> (x > 0))", data, expected);

  // Mark all arrays null. Expect all-null results.
  for (auto i = 0; i < data->size(); ++i) {
    data->childAt(0)->set_null(i, true);
  }

  expected = make_all_null_flat_vector<int32_t>(data->size());
  verify("find_first(c0, x -> (x > 0))", data, expected);
}

// Verify evaluation on a subset of rows in ArrayVector with only null and empty
// arrays..
TEST_F(FindFirstTest, emptyArrays) {
  auto data = make_row_vector({
      make_array_vector_from_json<int32_t>({
          "[]",
          "[]",
          "null",
          "[]",
          "[]",
          "null",
      }),
  });

  // Evaluate on all rows, but the first one.
  SelectivityVector rows(data->size());
  rows.setValid(0, false);
  rows.updateBounds();

  auto result = evaluate("find_first(c0, x -> (x > 0))", data, rows);
  auto expected = make_null_constant(TypeKind::INTEGER, data->size());
  pollux::test::assertEqualVectors(expected, result);
}

} // namespace
} // namespace kumo::pollux::functions
